import faiss
import pickle
from models.sentence_transformer import generate_embeddings
from models.t5_model import T5Model
from config import DATA_PATH, PREPROCESSED_PATH, EMBEDDINGS_PATH, INDEX_PATH
import torch
import os
import gc

def load_preprocessed_data():
    with open(PREPROCESSED_PATH, 'rb') as f:
        return pickle.load(f)

def load_embeddings():
    with open(EMBEDDINGS_PATH, 'rb') as f:
        return pickle.load(f)

def load_faiss_index():
    return faiss.read_index(INDEX_PATH)

def search_index(query, index, preprocessed_data, top_k=3):
    try:
        with torch.no_grad():
            query_embedding = generate_embeddings([query])
            if len(query_embedding) == 0:
                return []
                
            query_embedding = query_embedding[0].astype('float32')
            query_embedding = query_embedding.reshape(1, -1)
            
            distances, indices = index.search(query_embedding, top_k)
            return [preprocessed_data[i] for i in indices[0]]
            
    except Exception as e:
        print(f"搜索索引时出错: {str(e)}")
        return []

def cleanup():
    """清理资源"""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()

def main():
    try:
        # 设置为CPU模式
        os.environ['CUDA_VISIBLE_DEVICES'] = ''
        torch.set_num_threads(1)
        
        print("加载数据...")
        preprocessed_data = load_preprocessed_data()
        faiss_index = load_faiss_index()
        
        print("初始化模型...")
        t5_model = T5Model()
        
        while True:
            try:
                query = input("\n请输入查询（输入 'q' 退出）：")
                if query.lower() == 'q':
                    break
                
                print("搜索相关文档...")
                context_docs = search_index(query, faiss_index, preprocessed_data, top_k=3)
                
                if not context_docs:
                    print("未找到相关文档")
                    continue
                    
                print("生成答案...")
                context = " ".join(context_docs)
                answer = t5_model.generate_answer(query, context)
                
                print("\n答案:", answer)
                
            except Exception as e:
                print(f"处理查询时出错: {str(e)}")
                cleanup()
                
    except Exception as e:
        print(f"程序执行出错: {str(e)}")
    finally:
        if 't5_model' in locals():
            del t5_model
        cleanup()

if __name__ == "__main__":
    main()
